import asyncio
import json
import logging
from asyncio import Event
from json import JSONDecodeError
from datetime import datetime

from fastapi import APIRouter, BackgroundTasks
from pydantic import BaseModel
from starlette.websockets import WebSocket, WebSocketDisconnect
from websockets.exceptions import ConnectionClosedOK

from src.agents.base_qa import a_qa
from src.agents.file_qa import file_qa
from src.agents.load_document import storage_api, delete_files, delete_konwledge
from src.agents.prompt import summary_prompt, diff_prompt
from src.agents.query import query
from src.tools.dms_operator import dmsoperator
from websocket_schemas import ChatResponse

stop_event = Event() # 对话中断事件
# from src.correct.tasks import run_task
router = APIRouter(
    prefix="/agent",
    tags=["agent"],
    responses={404: {"description": "Not found"}},
)
class Integration(BaseModel):
    '''数据集成'''
    file_name:str
    file_id:str
    file_path:str
    knowledge_id:str
    knowledge_name:str
    operation_type: int
    parent_knowledge_id: str
    # parent_knowledge_name: str
    belong_file_id: str
# class CorrectBody(BaseModel):
#     before_id: str
#     current_id: str

# 数据集成；知识库数据入库
@router.post("/integration")
async def integration(integration_b:Integration):
    if integration_b.operation_type==0: #插入
        metadata = {"filename":integration_b.file_name,"file_id":integration_b.file_id,
                    "knowledge_id":integration_b.knowledge_id,"knowledge_name":integration_b.knowledge_name,
                    "parent_knowledge_id":integration_b.parent_knowledge_id,
                    "belong_file_id": integration_b.belong_file_id, "create_time": datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
        print(f"插入：{metadata}")
        delete_files(integration_b.belong_file_id)
        if integration_b.file_name.lower().endswith(('.doc', '.docx', 'txt')):
            contents =dmsoperator.get_fileconversion(integration_b.file_path,4)
            if contents is None:
                return {"status": "error", "message": f"{integration_b.file_name}文件转换失败"}
        else:
            contents =dmsoperator.get_textfile(integration_b.file_path)
        # contents = dmsoperator.get_textfile(integration_b.file_path)
        storage_api(file_contents=contents,metadata=metadata)
    elif integration_b.operation_type==1: # 删除文件
        print(f'删除文件：文件：{integration_b.file_name}，知识库：{integration_b.knowledge_name}')
        delete_files(integration_b.belong_file_id)
    elif integration_b.operation_type==2: # 修改
        # 先删除原有索引，再插入新的内容
        metadata = {"filename": integration_b.file_name, "file_id": integration_b.file_id,
                    "knowledge_id": integration_b.knowledge_id, "knowledge_name": integration_b.knowledge_name,
                    "parent_knowledge_id": integration_b.parent_knowledge_id,
                    "belong_file_id": integration_b.belong_file_id,
                    "create_time": datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
        print(f"修改：{metadata}")
        delete_files(integration_b.belong_file_id)
        if integration_b.file_name.lower().endswith(('.doc', '.docx', 'txt')):
            contents = dmsoperator.get_fileconversion(integration_b.file_path, 4)
            if contents is None:
                return {"status": "error", "message": f"{integration_b.file_name}文件转换失败"}
        else:
            contents = dmsoperator.get_textfile(integration_b.file_path)
        # contents = dmsoperator.get_textfile(integration_b.file_path)
        storage_api(file_contents=contents, metadata=metadata)
    elif integration_b.operation_type==3: # 删除知识库
        print(f'删除知识库：文件：{integration_b.file_name}，知识库：{integration_b.knowledge_name}')
        delete_konwledge(integration_b.parent_knowledge_id)
    return {"status": "success", "message": "数据入库成功"}


@router.websocket("/ws/qa")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    global stop_event # 这里做一个事件用于监听停止状态，但是仅仅限于在大模型的吐字的循环中中断；如果前端要做到在embedding或者大模型做长文本推理过程中还未吐字时候中断；那么让前端直接断开连接重连可以解决这个问题
    stop_event.clear()  # 确保事件是未触发状态
    while True:
        try:
            # Receive and send back the client message
            msg_text = await websocket.receive_text()
            # body_dict = json.loads(json_msg)
            # user_msg = body_dict["msg"]
            # req_type = body_dict["type"]
            # user_id = body_dict["user_id"]
            # if req_type == "base_qa":
            #     pass
            if msg_text == "stopAnswering" and not stop_event.is_set():
                stop_event.set() # 设置中断事件
                continue
            if msg_text == "stopAnswering" and stop_event.is_set():
                stop_event.clear()#  如果已经停止；清除中断事件<防止不停乱点>
                continue
            print(f"前端信息：{msg_text}")
            msg_json = json.loads(msg_text)
            user_msg = msg_json["msg"]
            qa_type = msg_json["type"]
            user_id = msg_json.get("user_id")
            knowledge_base_ids = msg_json.get("knowledge_base_ids")
            history = msg_json.get("history")
            file_ids = msg_json.get("file_ids")
            resp = ChatResponse(sender="human", message=user_msg, type="stream")
            await websocket.send_json(resp.dict())
            # Construct a response
            start_resp = ChatResponse(sender="bot", message="", type="start")
            await websocket.send_json(start_resp.dict())
            if not history:
                history = []
            if qa_type == "base_qa":
                question = user_msg
                #变成后台输出，不阻塞下一个请求；以便接受停止信号
                loop = asyncio.get_event_loop()
                loop.create_task(a_qa(websocket, question, history, stop_event))
                # await a_qa(websocket, question, history, stop_event)
            elif qa_type == "kg_qa":
                # 发送聊天信息、
                loop = asyncio.get_event_loop()
                loop.create_task(query(websocket,user_msg,knowledge_base_ids, user_id, stop_event))
                # await query(websocket,user_msg,[], stop_event) # 知识库qa
            elif qa_type == "summary":
                loop = asyncio.get_event_loop()
                if history and len(history)>0:
                    loop.create_task(file_qa(websocket, user_msg, history, file_ids,stop_event))
                else:
                    loop.create_task(file_qa(websocket, summary_prompt, history, file_ids,stop_event))
            elif qa_type == "diff_analyze":
                loop = asyncio.get_event_loop()
                if history and len(history)>0:
                    loop.create_task(file_qa(websocket, user_msg, history, file_ids,stop_event))
                else:
                    loop.create_task(file_qa(websocket, diff_prompt, history, file_ids,stop_event))
        except WebSocketDisconnect:
            logging.info("WebSocketDisconnect")
            # TODO try to reconnect with back-off
            break
        except ConnectionClosedOK:
            logging.info("ConnectionClosedOK")
            # TODO handle this?
            break
        except JSONDecodeError:
            logging.info("json decode error")
        # except Exception as e:
        #     logging.error(e)
        #     resp = ChatResponse(
        #         sender="bot",
        #         message="后台出现错误，请稍后再试，如果无法解决，请反馈给管理员。",
        #         type="error",
        #     )
        #     await websocket.send_json(resp.dict())